"""
测试分类搜索功能
"""

from db_utils import Database, VectorStore, DocumentManager, init_db


# 初始化数据库表结构
init_db()

# 创建数据库连接和向量存储实例
db = Database()
session = db.connect()  # 保存session对象用于后续关闭
vector_store = VectorStore(db)

# 创建文档管理器
doc_manager = DocumentManager(db)

print("测试分类搜索功能...")
print("====================")

# 测试搜索所有文档（不指定分类）
print("\n1. 搜索所有文档（不指定分类），关键词【西游记】:")
results_all = vector_store.search_similar_documents("西游记", top_k=10)
print(f"找到 {len(results_all)} 个结果")
for i, (doc_id, content, title, similarity, category_ids) in enumerate(results_all):
    print(
        f"  {i+1}. ID: {doc_id}, 标题: {title}, 相似度: {similarity:.4f}, 分类: {category_ids}"
    )

# 测试搜索分类1（中国古代小说）下的文档
print("\n2. 搜索分类1（中国古代小说）下的文档:")
results_category_1 = vector_store.search_similar_documents(
    "西游记", top_k=10, category_id=1
)
print(f"找到 {len(results_category_1)} 个结果")
for i, (doc_id, content, title, similarity, category_ids) in enumerate(
    results_category_1
):
    print(
        f"  {i+1}. ID: {doc_id}, 标题: {title}, 相似度: {similarity:.4f}, 分类: {category_ids}"
    )

# 测试搜索分类3（数学）下的文档
print("\n3. 搜索分类3（数学）下的文档:")
results_category_3 = vector_store.search_similar_documents(
    "三角函数", top_k=10, category_id=3
)
print(f"找到 {len(results_category_3)} 个结果")
for i, (doc_id, content, title, similarity, category_ids) in enumerate(
    results_category_3
):
    print(
        f"  {i+1}. ID: {doc_id}, 标题: {title}, 相似度: {similarity:.4f}, 分类: {category_ids}"
    )

# 测试搜索分类5（考研408）下的文档
print("\n4. 搜索分类5（考研408）下的文档:")
results_category_5 = vector_store.search_similar_documents(
    "C语言", top_k=10, category_id=5
)
print(f"找到 {len(results_category_5)} 个结果")
for i, (doc_id, content, title, similarity, category_ids) in enumerate(
    results_category_5
):
    print(
        f"  {i+1}. ID: {doc_id}, 标题: {title}, 相似度: {similarity:.4f}, 分类: {category_ids}"
    )

# 清理数据库连接
db.disconnect(session)

print("\n测试完成！")
